from com.demo1.ai.qianfan import demo_qf_ai_base_config as qfConfig
import requests
import json


def no_stream_v1(user_msg):
    url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/yi_34b_chat?access_token=" + qfConfig.get_access_token()

    payload = json.dumps({
        "messages": [
            # {
            #     "role": "assistant",
            #     "content": "你是百度提供的人工智能助手，你更擅长中文对话。你会为用户提供安全，有帮助，准确的回答。同时，你会拒绝一切涉及恐怖主义，种族歧视，黄色暴力等问题的回答。"
            # },
            # {
            #     "role": "assistant",
            #     "content": "1+1等于"
            # },
            {
                "role": "user",
                "content": user_msg
            }
        ]
    })
    headers = {
        'Content-Type': 'application/json'
    }

    response = requests.request("POST", url, headers=headers, data=payload)

    # print(response.text)
    return json.loads(response.text)["result"]


def stream_v1(user_msg):
    url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/yi_34b_chat?access_token=" + qfConfig.get_access_token()

    payload = json.dumps({
        "messages": [
            {
                "role": "user",
                "content": user_msg
            }
        ],
        "stream": True
    })
    headers = {
        'Content-Type': 'application/json'
    }

    response = requests.request("POST", url, headers=headers, data=payload, stream=True)
    # for line in response.iter_lines():
    #     if line:
    #         print(json.loads(line.decode("UTF-8").split("data: ")[1])["result"])
    return response


def no_stream_v2():

    url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/eb-instant?access_token=" + qfConfig.get_access_token()

    payload = json.dumps({
        "messages": [
            {
                "role": "user",
                "content": "你好"
            }
        ]
    })
    headers = {
        'Content-Type': 'application/json'
    }

    response = requests.request("POST", url, headers=headers, data=payload)

    print(response.headers)
    print(response.text)

if __name__ == '__main__':
#     stream_v1("中国GDP前十的省份")
#     # no_stream_v1("1+1")
    no_stream_v2()
